import numpy as np

# 朴素贝叶斯算法训练函数
# 也就是从一组已标记（明确了侮辱性和非侮辱性）数据中计算出：侮辱性文档中每个单词出现的频率，
# 非侮辱性文档中每个单词出现的频率，所有文档中侮辱性文档频率。

def trainNB0(trainMatrix, trainCategory):
    '''
    参数：
        trainMatrix: 训练数据集（矩阵），其中每一行就是一个词向量
        trainCategory: 分类标签集
    返回值：
        p0V: 非侮辱性文档中，词汇表中每个单词出现的频率
        p1V: 侮辱性文档中，词汇表中每个单词出现的频率
        pAbusive: 所有文档中侮辱性文档出现的频率
    '''
    numDocs = len(trainMatrix) # 训练集中所有文档数目
    numWords = len(trainMatrix[0]) # 词向量长度
    pAbusive = sum(trainCategory)/float(numDocs) # 所有文档中，侮辱性文档频率

    # 初始化：侮辱性文档中每个单词出现与否（次数）；非侮辱性文档中每个单词出现与否（次数）
    # p0Num = np.zeros(numWords); p1Num = np.zeros(numWords)
    p0Num = np.ones(numWords); p1Num = np.ones(numWords)
    # 初始化：侮辱性文档中出现的单词总数；非侮辱性文档中出现的单词总数
    # p0Denom = 0.0; p1Denom = 0.0
    p0Denom = 2.0; p1Denom = 2.0

    for i in range(numDocs):
        if trainCategory[i]==1:
            p1Num += trainMatrix[i]
            p1Denom += sum(trainMatrix[i])
        else:
            p0Num += trainMatrix[i]
            p0Denom += sum(trainMatrix[i])
    # 防止出现下溢
    # p1V = p1Num/p1Denom
    # p0V = p0Num/p0Denom
    p1V = np.log(p1Num/p1Denom)
    p0V = np.log(p0Num/p0Denom)
    return p0V, p1V, pAbusive

# 朴素贝叶斯分类函数
def classifyNB(vec2Classify, p0Vec, p1Vec,pClass1):
    # vec2Classify*p1Vec就是取出对应的侮辱性文档中，vec2Classify中出现的词的概率
    rateOfWords = vec2Classify*p1Vec # 调试用
    p1 = sum(vec2Classify*p1Vec)+np.log(pClass1)
    p0 = sum(vec2Classify*p0Vec)+np.log(1-pClass1)
    if p1>p0:
        return 1
    else:
        return 0

# 辅助函数1：创建词汇表
# 就是将训练集中的所出现的每个单词构成一个构成一个无重复的向量

def createVocabList(dataSet):
    '''
    创建词汇表
    参数:
        dataSet: 二维列表，每一行就是一篇文档中出现的单词
    返回值:
        词汇表，也就是训练集中出现的所有单词的列表
    '''
    vocabList = set()
    for document in dataSet:
        vocabList = vocabList | set(document)
    return list(vocabList)

# 辅助函数2：文档向量化
# 也就是把每个文档中的单词，依据词汇表的位置标记为1，那么文档中没有出现的单词就是0

def setOfWords2Vec(vocabList, inputSet):
    '''
    params
        vocabList: 词汇表,也就是createVocabList()返回的训练集中所有出现的单词列表
        inputSet: 输入的一篇文档(这里已经是文档单词的列表)
    return
        returnVec: 长度为词汇表长度的向量, 其元素为0,1, 1代表词汇表中的某个词出现在该文档

    '''
    # 生成与词汇表长度一致的全0向量（列表）
    returnVec = [0]*len(vocabList)

    for word in inputSet:
        if word in vocabList:
            i = vocabList.index(word)
            returnVec[i] = 1
    return returnVec


if __name__ == '__main__':
    # 测试数据
    postingList=[
        ['my', 'dog', 'has', 'flea', 'problems', 'help', 'please'],
        ['maybe', 'not', 'take', 'him', 'to', 'dog', 'park', 'stupid'],
        ['my', 'dalmation', 'is', 'so', 'cute', 'I', 'love', 'him'],
        ['stop', 'posting', 'stupid', 'worthless', 'garbage'],
        ['mr', 'licks','ate', 'my', 'steak', 'how', 'to', 'stop', 'him'],
        ['quit', 'buying', 'worthless', 'dog', 'food', 'stupid']
    ]
    classVec=[0,1,0,1,0,1] #1 代表侮辱性文字, 0代表正常言论

    # 第一步：根据训练集建立词汇表
    vocabList = createVocabList(postingList)
    vocabList.sort() # 添加一步排序，以便固定下来
    # 将第一个文档向量化
    doc1Vec = setOfWords2Vec(vocabList,postingList[0])
    # print(doc1Vec)

    # 第二步：构建训练集
    trainMat = []
    for document in postingList:
        trainMat.append(setOfWords2Vec(vocabList,document))
    # print(trainMat)

    # 第三步：训练分类器
    p0v, p1v, pAbu = trainNB0(trainMat, classVec)
    
    # 第四步：使用分类器
    # testEntry = ['love', 'my', 'dalmation']
    testEntry = ['stupid', 'garbage']
    thisDoc=np.array(setOfWords2Vec(vocabList,testEntry))
    # print(thisDoc)
    ret = classifyNB(thisDoc,p0v, p1v,pAbu)
    print(ret)
